from typing import List
import torch

from hj_dataset_devkit import Frame

from model_frame import build_task_assigner, build_model, HJItem, DetectObject

from common_utils import read_yaml
from components.obj_detector import ObjDetector

__all__ = ['BEVGridDetector']

def load_state_to_model(pth_path: str, model: torch.nn.Module) -> None:
    checkpoint: dict = torch.load(pth_path, 'cpu')
    model.load_state_dict(checkpoint['model_state_dict'])

class BEVGridDetector(ObjDetector):
    def __init__(self, train_cfg_path: str, model_pth_path: str, device: str) -> None:
        super().__init__()
        train_cfg = read_yaml(train_cfg_path)
        self.__task = build_task_assigner(train_cfg['task'])
        self.__model = build_model(train_cfg['model'])
        load_state_to_model(model_pth_path, self.__model)
        print(f'INFO: detector model is loaded from {model_pth_path}')
        self.__device = torch.device(device)
        self.__model.to(self.__device)
        self.__model.eval()

    def process_frame(self, frame: Frame) -> List[DetectObject]:
        item = HJItem(frame.get_lidar_cloud('true_value')[:, :4], None)
        in_data = torch.from_numpy(self.__task.assign_input(item)).unsqueeze(0)
        with torch.no_grad():
            out_data = self.__model(in_data.to(self.__device))
        return self.__task.parse_model_output(out_data)[0]
